# -*- coding: utf-8 -*-
"""
带电粒子在二维电磁场中运动的可视化脚本 (V4.1 - 优化自定义场景)

功能:
- 通过 CONFIG 字典自定义所有参数。
- 支持多个、有边界、可重叠的电场和磁场。
- 可选择 Euler 法或更精确的 RK4 算法。
- 可视化电场矢量和磁场方向符号。
- 实时数据显示面板，可展示速度、加速度、受力等信息。
- 动画引擎优化，通过“每帧步数”解决轻粒子模拟的卡顿问题。
- 修正: 向内磁场符号正确显示为 ⊗。
- 优化: custom 场景配置为一个经典的“摆线运动”模型。
"""
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.patches import Rectangle

# ==============================================================================
# --- 用户配置区域 ---
# ==============================================================================

# 选择一个预带电粒子在二维电磁场中运动的可视化脚本 (V4.1 - 优化自定义场景)设场景进行模拟: 'cyclotron', 'accelerator', 'custom'
CURRENT_SCENARIO = 'custom' 

# --- 预设场景配置 ---
SCENARIO_PRESETS = {
    'cyclotron': {
        "particle": {
            "mass": 9.1e-31, "charge": -1.6e-19, # 电子
            "initial_position": np.array([0.0, -0.05]),
            "initial_velocity": np.array([1.0e7, 0.0]), # 速度提高以观察效果
        },
        "fields": {
            "electric": [],
            "magnetic": [{"value": 1.1e-3, "bounds": [-np.inf, -np.inf, np.inf, np.inf]}]
        },
        "simulation": {
            "use_gravity": False,
            "dt": 2e-12, # 极小的时间步长保证电子回旋的精度
            "steps_per_frame": 50 # 每一帧动画更新50个物理步，保证动画流畅
        },
        "visualization": {"plot_bounds": [-0.07, -0.07, 0.07, 0.07]}
    },
    'accelerator': {
        "particle": {
            "mass": 1.67e-27, "charge": 1.6e-19, # 质子
            "initial_position": np.array([-0.5, 0.0]),
            "initial_velocity": np.array([0.0, 0.0]),
        },
        "fields": {
            "electric": [{"vector": np.array([20.0, 0.0]), "bounds": [-np.inf, -np.inf, np.inf, np.inf]}],
            "magnetic": []
        },
        "simulation": {"use_gravity": False, "dt": 1e-8, "steps_per_frame": 10},
        "visualization": {"plot_bounds": [-0.5, -0.2, 0.5, 0.2]}
    },
    
    # --- 这里是为您完善后的 'custom' 场景 ---
    # 场景描述: 摆线 (Cycloid) 运动
    # 一个带正电的粒子在相互垂直的匀强电场（向下）和匀强磁场（向内）中从静止开始运动。
    # 粒子先在电场作用下向下加速，然后磁场力使其向右偏转，最终形成优美的摆线轨迹。
    'custom': {
        "particle": {
            "mass": 1.67e-27, "charge": 1.6e-19, # 质子
            "initial_position": np.array([0.0, 0.0]), # 从原点开始
            "initial_velocity": np.array([0.0, 0.0]), # 从静止开始
        },
        "fields": {
            "electric": [
                {
                    "vector": np.array([0.0, -0.5]), # 向下的匀强电场
                    "bounds": [-np.inf, -np.inf, np.inf, np.inf] # 全局有效
                }
            ],
            "magnetic": [
                {
                    "value": -1e-3, # 指向纸面内 (-z) 的匀强磁场
                    "bounds": [-np.inf, -np.inf, np.inf, np.inf] # 全局有效
                }
            ]
        },
        "simulation": {
            "use_gravity": False, # 忽略重力，观察纯粹的电磁效应
            "dt": 1e-8,
            "steps_per_frame": 15
        },
        "visualization": {
            "plot_bounds": [-0.01, -0.025, 0.07, 0.005] # 调整了视图以更好地容纳轨迹
        }
    }
}

# --- 全局配置 ---
CONFIG = {
    "particle": SCENARIO_PRESETS[CURRENT_SCENARIO]['particle'],
    "fields": SCENARIO_PRESETS[CURRENT_SCENARIO]['fields'],
    "simulation": {
        "total_time": 1e-4, # 总模拟物理时间
        "use_gravity": SCENARIO_PRESETS[CURRENT_SCENARIO]['simulation']['use_gravity'],
        "g": 9.81,
        "solver": 'rk4',
        "dt": SCENARIO_PRESETS[CURRENT_SCENARIO]['simulation']['dt'],
        "steps_per_frame": SCENARIO_PRESETS[CURRENT_SCENARIO]['simulation']['steps_per_frame']
    },
    "visualization": {
        "plot_bounds": SCENARIO_PRESETS[CURRENT_SCENARIO]['visualization']['plot_bounds'],
        "show_velocity_vector": True,
        "show_acceleration_vector": True,
        "vector_scale": 0.003, # 调整了矢量大小以适应新场景
        "trail_length": 200000,
        "show_field_vectors": True,
        "field_vector_density": 15,
        "show_info_panel": True,
        "info_force_breakdown": True,
    }
}

# ==============================================================================
# --- 模拟器核心代码 (保持不变) ---
# ==============================================================================

class ParticleSimulator:
    def __init__(self, config):
        self.config = config
        self.particle_conf = self.config["particle"]
        self.sim_conf = self.config["simulation"]
        self.vis_conf = self.config["visualization"]
        self.fields_conf = self.config["fields"]

        self.state = np.array([
            self.particle_conf["initial_position"][0], self.particle_conf["initial_position"][1],
            self.particle_conf["initial_velocity"][0], self.particle_conf["initial_velocity"][1]
        ])
        self.mass = self.particle_conf["mass"]
        self.charge = self.particle_conf["charge"]
        
        self.force_E, self.force_B, self.force_G, self.total_force, self.acceleration = [np.array([0.0, 0.0]) for _ in range(5)]

        self.history_x, self.history_y = [self.state[0]], [self.state[1]]
        self._setup_plot()

    def _is_in_bounds(self, pos, bounds):
        x_min, y_min, x_max, y_max = bounds
        return x_min <= pos[0] <= x_max and y_min <= pos[1] <= y_max

    def _calculate_forces(self, pos, vel):
        total_E, total_Bz = np.array([0.0, 0.0]), 0.0
        for field in self.fields_conf["electric"]:
            if self._is_in_bounds(pos, field["bounds"]): total_E += field["vector"]
        for field in self.fields_conf["magnetic"]:
            if self._is_in_bounds(pos, field["bounds"]): total_Bz += field["value"]
        
        self.force_E = self.charge * total_E
        self.force_B = self.charge * np.array([vel[1] * total_Bz, -vel[0] * total_Bz])
        self.force_G = np.array([0, -self.mass * self.sim_conf["g"]]) if self.sim_conf["use_gravity"] else np.array([0.0, 0.0])
        
        self.total_force = self.force_E + self.force_B + self.force_G

    def _derivatives(self, state):
        pos, vel = state[0:2], state[2:4]
        self._calculate_forces(pos, vel)
        self.acceleration = self.total_force / self.mass
        return np.array([vel[0], vel[1], self.acceleration[0], self.acceleration[1]])
    
    def _update_physics_step(self):
        """执行单步物理计算（Euler或RK4）"""
        dt = self.sim_conf["dt"]
        if self.sim_conf['solver'] == 'rk4':
            k1 = self._derivatives(self.state)
            k2 = self._derivatives(self.state + 0.5 * dt * k1)
            k3 = self._derivatives(self.state + 0.5 * dt * k2)
            k4 = self._derivatives(self.state + dt * k3)
            self.state += (dt / 6.0) * (k1 + 2*k2 + 2*k3 + k4)
        else:
            self.state += self._derivatives(self.state) * dt

    def _setup_plot(self):
        self.fig, self.ax = plt.subplots(figsize=(12, 10))
        bounds = self.vis_conf["plot_bounds"]
        self.ax.set(xlim=(bounds[0], bounds[2]), ylim=(bounds[1], bounds[3]),
                    xlabel="x (m)", ylabel="y (m)",
                    title=f"Particle Motion (Solver: {self.sim_conf['solver'].upper()})")
        self.ax.set_aspect('equal', adjustable='box')
        self.ax.grid(True, linestyle='--', alpha=0.6)
        
        if self.vis_conf["show_field_vectors"]: self._draw_fields()

        self.trail, = self.ax.plot([], [], 'b-', lw=1.5, label='Trajectory')
        self.particle, = self.ax.plot([], [], 'ro', markersize=8, label='Particle')
        
        self.vel_vector = self.ax.quiver([],[],[],[], color='g', scale=1, units='xy', angles='xy', scale_units='xy', label='Velocity' if self.vis_conf["show_velocity_vector"] else "")
        self.acc_vector = self.ax.quiver([],[],[],[], color='orange', scale=1, units='xy', angles='xy', scale_units='xy', label='Acceleration' if self.vis_conf["show_acceleration_vector"] else "")
        
        if self.vis_conf["show_info_panel"]:
            self.info_text = self.ax.text(0.02, 0.98, '', transform=self.ax.transAxes, fontsize=9, verticalalignment='top', bbox=dict(boxstyle='round,pad=0.5', fc='white', alpha=0.7))
        else:
            self.info_text = None 
        
        self.ax.legend(loc='upper right')

    def _draw_fields(self):
        bounds, density = self.vis_conf["plot_bounds"], self.vis_conf["field_vector_density"]
        x_grid, y_grid = np.linspace(bounds[0], bounds[2], density), np.linspace(bounds[1], bounds[3], density)
        handles, labels = self.ax.get_legend_handles_labels()

        for field in self.fields_conf["electric"]:
            Ex, Ey, X, Y = [], [], [], []
            for x in x_grid:
                for y in y_grid:
                    if self._is_in_bounds(np.array([x, y]), field["bounds"]):
                        X.append(x); Y.append(y); Ex.append(field["vector"][0]); Ey.append(field["vector"][1])
            if X:
                label = 'E-Field' if 'E-Field' not in labels else None
                self.ax.quiver(X, Y, Ex, Ey, color='cyan', alpha=0.7, headwidth=2, headlength=3, label=label)
                if label: labels.append(label)

        for field in self.fields_conf["magnetic"]:
            X_in, Y_in, X_out, Y_out = [], [], [], []
            for x in x_grid:
                for y in y_grid:
                    if self._is_in_bounds(np.array([x, y]), field["bounds"]):
                        if field["value"] > 0: X_out.append(x); Y_out.append(y)
                        elif field["value"] < 0: X_in.append(x); Y_in.append(y)
            if X_out:
                label_out = 'B-Field (out, ⊙)' if 'B-Field (out, ⊙)' not in labels else None
                self.ax.scatter(X_out, Y_out, marker='o', s=50, facecolors='none', edgecolors='m', alpha=0.7, label=label_out)
                self.ax.scatter(X_out, Y_out, marker='.', s=5, color='m', alpha=0.7)
                if label_out: labels.append(label_out)
            if X_in:
                label_in = 'B-Field (in, ⊗)' if 'B-Field (in, ⊗)' not in labels else None
                self.ax.scatter(X_in, Y_in, marker='o', s=30, facecolors='none', edgecolors='m', alpha=0.7, label=label_in)
                self.ax.scatter(X_in, Y_in, marker='x', s=20, color='m', alpha=0.7)
                if label_in: labels.append(label_in)

    def _update_info_panel(self):
        pos, vel = self.state[0:2], self.state[2:4]
        text = [f"Position:     ({pos[0]:.2e}, {pos[1]:.2e}) m",
                f"Velocity:     ({vel[0]:.2e}, {vel[1]:.2e}) m/s",
                f"Speed:        {np.linalg.norm(vel):.2e} m/s",
                f"Acceleration: ({self.acceleration[0]:.2e}, {self.acceleration[1]:.2e}) m/s²"]
        if self.vis_conf["info_force_breakdown"]:
            text.extend(["--- Forces (N) ---",
                         f"F_Electric: ({self.force_E[0]:.2e}, {self.force_E[1]:.2e})",
                         f"F_Magnetic: ({self.force_B[0]:.2e}, {self.force_B[1]:.2e})",
                         f"F_Gravity:  ({self.force_G[0]:.2e}, {self.force_G[1]:.2e})",
                         f"F_Total:    ({self.total_force[0]:.2e}, {self.total_force[1]:.2e})"])
        else:
            text.append(f"Total Force:  ({self.total_force[0]:.2e}, {self.total_force[1]:.2e}) N")
        self.info_text.set_text("\n".join(text))

    def _update_frame(self, frame):
        for _ in range(self.sim_conf["steps_per_frame"]):
            self._update_physics_step()
            self.history_x.append(self.state[0])
            self.history_y.append(self.state[1])
        
        trail_len = self.vis_conf["trail_length"]
        if len(self.history_x) > trail_len:
            self.history_x = self.history_x[-trail_len:]
            self.history_y = self.history_y[-trail_len:]

        self.trail.set_data(self.history_x, self.history_y)
        self.particle.set_data([self.state[0]], [self.state[1]])
        
        pos, vel = self.state[0:2], self.state[2:4]
        if self.vis_conf["show_velocity_vector"]:
            v_norm = np.linalg.norm(vel)
            if v_norm > 1e-9:
                scaled_v = vel / v_norm * self.vis_conf["vector_scale"]
                self.vel_vector.set_offsets(pos); self.vel_vector.set_UVC(scaled_v[0], scaled_v[1])
        if self.vis_conf["show_acceleration_vector"]:
            a_norm = np.linalg.norm(self.acceleration)
            if a_norm > 1e-9:
                scaled_a = self.acceleration / a_norm * self.vis_conf["vector_scale"]
                self.acc_vector.set_offsets(pos); self.acc_vector.set_UVC(scaled_a[0], scaled_a[1])

        if self.info_text: self._update_info_panel()

        artists = [self.trail, self.particle, self.vel_vector, self.acc_vector]
        if self.info_text: artists.append(self.info_text)
        return artists

    def run_animation(self):
        num_frames = int(self.sim_conf["total_time"] / (self.sim_conf["dt"] * self.sim_conf["steps_per_frame"]))
        ani = animation.FuncAnimation(self.fig, self._update_frame, frames=num_frames,
                                      interval=20, blit=True, repeat=False)
        plt.show()

# ==============================================================================
# --- 脚本主入口 ---
# ==============================================================================
if __name__ == '__main__':
    simulator = ParticleSimulator(CONFIG)
    simulator.run_animation()